rm(list = ls())
library(reshape2)
library(ggplot2)
library(magrittr)
library(stringr)
devtools::load_all()

# 参数设置
vlast <- rep(0,100)
k0 <- seq(0.06,6,0.06)
beta <- .98
delta <-  .1
theta <- .36
numits <- 240

# 图形初始化
picdata <- data.frame(k0 = k0, v = vlast)

# 值函数迭代numits次。在每个k=(0.01，6.2)上以0作为初值开始迭代
v <- k_opt <- numeric(length(k0))
for (i in 1:numits){
  # 寻找每个k点的最优值函数
  for (j in 1:length(k0)){
    # 优化
    ans <- optimize(valfun, interval = c(0.01,6.2), kt = k0[j], beta = beta, theta = theta,
             delta = delta, k0 = k0, vlast = vlast)

    v[j] <- -ans$objective
    k_opt[j] <- ans$minimum
  }

  # 每48次迭代存储一下
  if (i %% 48 == 0){
    print(i)
    picdata[,paste('v',as.character(i), sep = '')] <- v
    picdata[,paste('k_opt',as.character(i), sep = '')] <- k_opt
  }

  # 替换上一次的值函数
  vlast <- v
}

# value function
ans <- melt(picdata, id.vars = 'k0')
ggplot(ans[str_detect(ans$variable,'^v[1-9]'),], aes(x = k0, y = value, color = variable)) +
  geom_line() + theme_bw()

# policy function
ggplot(ans[str_detect(ans$variable,'k_opt240'),], aes(x = k0, y = value, color = variable)) +
  geom_line() + theme_bw()

# consume
picdata$k_opt240^theta+(1-delta)*picdata$k_opt240-lead(picdata$k_opt240)

